#!/usr/bin/env python
"""
rankAbundance.py

Description:-

Give one or more pickle files as input and this shall output a white space
delimited table showing calculated biodiversity information. This includes
Shannon, SImpson & JackKnife Indices, which are indices which can be
interpreted as representations of species evenness, richness, and sampling
bias, repsectively.

Some useful options included with this program:-

    -collapse-at-rank <rank>   :-
            Collapse the result tree at rank <rank>.
    -groupX [group1files] [...] [-group2 ... [...]] :-
            * used to combine and average diversity indices over
            multiple results files. 


"""

from ssummolib import dict_walk, Options, collapse_at_rank, combine_dicts
from math import log,fsum
import cPickle as pickle
import sys

class LocalOptions( Options ):
    def __init__( self, args ):
        Options.__init__(self)
        self.options = {'-in': [],
                           '-out' : sys.stdout,
                           '-collapse-at-rank' : None,
                           }
        self.singleargs = ['-out','-collapse-at-rank']
        self.multiargs = ['-in']
        self.helpTxt = {'-in' : 'Input pickle file. Required.',
                           '-out': 'Output table. [Standard output]',
                           '-collapse-at-rank' : 'Desired rank at which to collapse the results tree',
                          }
        self.useage = 'python {0} [options] <results file>.pkl'.format( __file__ )
        self.parseArgs( args )

class Indices:
    def __init__( self, options ):
        self.options = options
    def tally_out( self , combined=False ):
        """Writes rank abundance information to options['-out'],
        and returns the counts, probabilities and total number of
        taxa.
        """
        counts = []
        for path, node in dict_walk('', self.dictionary ):
            if 'accessions' in node:
                if combined:
                    count = 0
                    for dataset in node['accessions']:
                        for prob in dataset:
                            if float( prob ) > 0.:
                                count += 1
                    counts.append( count )
                else:
                    counts.append( len( node['accessions'] ) )
        if self.options['-out'] != sys.stdout and type(self.options['-out']) is not file:
            self.options['-out'] = file( self.options['-out'] ,'w' )
        counts = sorted(counts, reverse=True )
        total = sum(counts )
        probs = [ float(count) / total for count in counts ]
        self.options['-out'].write( '\n'.join( [ str(100.*prob) for prob in probs ] ))
        self.options['-out'].flush()
        return counts, probs, total

    def ShannonIndex( self, probs ):
        """Give the probabilities as returned from tally_out() and total
        species count. This returns the tuple H_prime, H_max
        """
        H_prime = - fsum( [ p * log(p) for p in probs ] )
        H_max = log( len(probs) )
        return ( H_prime, H_max ,)

    def ShannonError( self, probs, counts, total ):
        pass

    def SimpsonsIndex( self, counts, total ):
        """Simpson's index of Diversity.
        Returns a value from 0 - 1, where 0 shows no
        diversity, and 1 shows infinite diversity.
        """
        diversity = 1. - ( fsum( [ n * (n-1.) for n in counts ] )  / (total * (total -1. ))   )
        return diversity

    def jackKnife( self, counts, total ):
        k = counts.count(1)  # number of 'unique' species.
        n = len(counts)
        S = n + ( (n-1.)/n)**k  # jackknife estimate.
        return S

    def load( self , results_dictionary , combined=False ):
        """Load the dictionary. If the option '-collapse-at-rank' is given,
        then we shall also collapse the dictionary at the specificied rank.
        Results dictionary is not returned, but is stored as:-
        self.dictionary
        """
        if self.options['-collapse-at-rank'] is not None:
            results_dictionary = collapse_at_rank( results_dictionary , self.options , combined )
        self.dictionary = results_dictionary

def main( options ):
    print '#{0} {1} {2} {3} {4}'.format( ' File'.ljust(25),
                                         'Shannon'.ljust(20) ,
                                         'Max Evenness'.ljust(25),
                                         'Simpsons'.ljust(25),
                                         'JackKnife'.ljust(25),
                                         )
    print '#{0} {1} {2} {3} {4}'.format( ' name'.ljust(25),
                                         '(H`)'.ljust(20) ,
                                         '(Hmax)'.ljust(25),
                                         'Richness'.ljust(25),
                                         '(sampling bias)'.ljust(25),
                                     )
    in_files = []
    if len( options.regopts ) > 0:
        groups = [ group for group in sorted( options.regopts.keys() )]
        dictionaries = []
        for group in groups:
            results = []
            for infile in options.regopts[group]:
                with file( infile ,'rb' ) as handle:
                    results.append( pickle.load( handle ) )
            combined_dict = combine_dicts( results )
            dictionaries.append( combined_dict )
            in_files.append( 'GROUP{0}'.format( group ))
        combined = True
    else:
        combined = False
        dictionaries = []
        for in_file in options['-in']:
            with file( in_file , 'rb' ) as inFile:
                dictionary = pickle.load( inFile )
            dictionaries.append( dictionary )
            in_files.append( in_file )
        
    for i, in_file in enumerate( in_files ):
        calculator = Indices( options )
        calculator.load( dictionaries[i], combined )
        counts, probs,total = calculator.tally_out( )
        richness, max_evenness = calculator.ShannonIndex( probs )
        diversity = calculator.SimpsonsIndex( counts, total )
        jack = calculator.jackKnife( counts, total ) 
        print '{0}  {1} {2} {3} {4}'.format( in_file,
                                         repr( richness    ).ljust(20) ,
                                         repr( max_evenness).ljust(25) ,
                                         repr( diversity   ).ljust(25) ,
                                         repr( jack        ).ljust(25) )
    if len( options['-in'] ) > 0:
        pass

if __name__ == '__main__':
    args = sys.argv[1:]
    myoptions = LocalOptions( args )
    main( myoptions )
